import torch
from torchvision import datasets, transforms

def load_mnist(batch_size=64):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_set = datasets.MNIST(
        './data', train=True, download=True, transform=transform)
    test_set = datasets.MNIST(
        './data', train=False, transform=transform)

    train_loader = torch.utils.data.DataLoader(
        train_set, batch_size=batch_size, shuffle=True, pin_memory=True)
    test_loader = torch.utils.data.DataLoader(
        test_set, batch_size=1000, shuffle=False)

    return train_loader, test_loader